#


"""
stock_base
"""
import uuid
import datetime
import json
import os
import copy
import threading
import requests
import pandas as pd
import pymongo
from qaenv import (eventmq_ip, eventmq_password, eventmq_port,
                   eventmq_username, mongo_ip)

import QUANTAXIS as QA
from QUANTAXIS.QAData.data_resample import QA_data_min_resample, QA_data_min_to_day
from QUANTAXIS.QAARP import QA_Risk, QA_User
from QUANTAXIS.QAEngine.QAThreadEngine import QA_Thread
from QUANTAXIS.QAUtil.QAParameter import MARKET_TYPE, RUNNING_ENVIRONMENT, ORDER_DIRECTION
from QAPUBSUB.QAPUBSUB.consumer import subscriber_topic,  subscriber
from QAPUBSUB.QAPUBSUB.producer import publisher_routing
from QAStrategy.QAStrategy.qa_muti_freq_base import QAStrategyMutiFreqenceBase
from QIFIAccount import QIFI_Account
from QIFIAccount.QIFIAccount.QARealtimeStockSim import QIFI_StockSIM_Account


class QAStrategyMutiFreqStockBase(QAStrategyMutiFreqenceBase):

    def __init__(self, username, password, code=['000001'], frequences=['1min', 'day'], strategy_id='QA_STRATEGY',
                 risk_check_gap=1, portfolio='default', init_cash=1000000,
                 start='2019-01-01', end='2019-10-21', send_wx=False, taskid=None, mongo_ip=mongo_ip):
        super().__init__(username=username, password=password, code=code, frequences=frequences, strategy_id=strategy_id,
                         risk_check_gap=risk_check_gap, portfolio=portfolio, start=start, end=end, send_wx=send_wx,
                         data_host=eventmq_ip, data_port=eventmq_port, data_user=eventmq_username, data_password=eventmq_password,
                         trade_host=eventmq_ip, trade_port=eventmq_port, trade_user=eventmq_username, trade_password=eventmq_password,
                         taskid=taskid, mongo_ip=mongo_ip, init_cash=init_cash)

        self.code = code
        self.send_wx = send_wx

    def subscribe_datas(self, code, frequences, data_host, data_port, data_user, data_password, model='py'):
        for freq in frequences:
            if freq == '1min':
                self.sub_1min = self.subscribe_data(code, freq, data_host, data_port, data_user, data_password, model)
                self.sub_1min.callback = self.callback_1min
            elif freq == '5min':
                self.sub_5min = self.subscribe_data(code, freq, data_host, data_port, data_user, data_password, model)
                self.sub_5min.callback = self.callback_5min
            elif freq == '15min':
                self.sub_15min = self.subscribe_data(code, freq, data_host, data_port, data_user, data_password, model)
                self.sub_15min.callback = self.callback_15min
            elif freq == '30min':
                self.sub_30min = self.subscribe_data(code, freq, data_host, data_port, data_user, data_password, model)
                self.sub_30min.callback = self.callback_30min
            elif freq == '60min':
                self.sub_60min = self.subscribe_data(code, freq, data_host, data_port, data_user, data_password, model)
                self.sub_60min.callback = self.callback_60min
            elif freq == 'day':
                self.sub_day = self.subscribe_data(code, freq, data_host, data_port, data_user, data_password, model)
                self.sub_day.callback = self.callback_day
            elif freq == 'week':
                self.sub_week = self.subscribe_data(code, freq, data_host, data_port, data_user, data_password, model)
                self.sub_week.callback = self.callback_week
            elif freq == 'month':
                self.sub_month = self.subscribe_data(code, freq, data_host, data_port, data_user, data_password, model)
                self.sub_month.callback = self.callback_month
            elif freq == 'year':
                self.sub_year = self.subscribe_data(code, freq, data_host, data_port, data_user, data_password, model)
                self.sub_year.callback = self.callback_year
            else:
                print('Unknown frequency %s' % freq)


    def subscribe_data(self, code, frequence, data_host, data_port, data_user, data_password, model='py'):
        """[summary]

        Arguments:
            code {[type]} -- [description]
            frequence {[type]} -- [description]
        """
        if self.get_freq_type(frequence) == QA.FREQUENCE.REALTIME_MIN:
            freq = frequence[:-3]
            print('current freq is %s min' % freq)
            sub = subscriber_topic(exchange='realtime_stock_{}_min'.format(freq),
                                        host=data_host,
                                        port=data_port,
                                        user=data_user,
                                        password=data_password,
                                        routing_key='stock')
        elif self.get_freq_type(frequence) == QA.FREQUENCE.DAY:
            freq = frequence[:-3]
            print('current freq is %s day' % freq)
            # TODO 暂时不用resample
            sub = subscriber(host=eventmq_ip, exchange='realtime_stock_day')
        elif self.get_freq_type(frequence) == QA.FREQUENCE.WEEK:
            sub = None
        elif self.get_freq_type(frequence) == QA.FREQUENCE.MONTH:
            sub = None
        elif self.get_freq_type(frequence) == QA.FREQUENCE.YEAR:
            sub = None
        else:
            sub = None
        # for item in code:
        #     self.sub.add_sub(exchange='realtime_stock_{}_min'.format(freq), routing_key=item)
        # self.sub.callback = self.callback
        return sub

    def upcoming_data(self, new_bar):
        """upcoming_bar :

        Arguments:
            new_bar {json} -- [description]
        """
        self.update_account()
        # self.positions.on_price_change(float(new_bar['close']))
        # 这里只取最后一个bar
        self.on_bar(new_bar.loc[(new_bar.index.levels[0][-1], slice(None)), :])
        # 全部bar
        # self.on_bar(new_bar)

    def merge_data(self, old_data, new_data):
        tmp = old_data
        inter = tmp.index.levels[0].intersection(new_data.index.levels[0])
        if len(inter) > 0:
            tmp = tmp.drop(index=inter)
        return tmp.append(new_data)

    def callback(self, a, b, c, body):
        self.new_data = pd.read_msgpack(body)
        bar = self.new_data.rename(columns=lambda x: x.replace('vol', 'volume'))
        self.upcoming_data(bar)

    def callback_1min(self, a, b, c, body):
        new_data = pd.read_msgpack(body)
        bar = new_data.rename(columns=lambda x: x.replace('vol', 'volume'))
        self._market_data_1min = self.merge_data(self._market_data_1min, bar)
        # 这里只有1min触发on_bar，所以如果使用多周期策略，必须包含1min周期
        self.upcoming_data(bar)

    def callback_5min(self, a, b, c, body):
        new_data = pd.read_msgpack(body)
        bar = new_data.rename(columns=lambda x: x.replace('vol', 'volume'))
        self._market_data_5min = self.merge_data(self._market_data_5min, bar)

    def callback_15min(self, a, b, c, body):
        new_data = pd.read_msgpack(body)
        bar = new_data.rename(columns=lambda x: x.replace('vol', 'volume'))
        self._market_data_15min = self.merge_data(self._market_data_15min, bar)

    def callback_30min(self, a, b, c, body):
        new_data = pd.read_msgpack(body)
        bar = new_data.rename(columns=lambda x: x.replace('vol', 'volume'))
        self._market_data_30min = self.merge_data(self._market_data_30min, bar)

    def callback_60min(self, a, b, c, body):
        new_data = pd.read_msgpack(body)
        bar = new_data.rename(columns=lambda x: x.replace('vol', 'volume'))
        self._market_data_60min = self.merge_data(self._market_data_60min, bar)

    def callback_day(self, a, b, c, body):
        new_data = pd.read_msgpack(body)
        bar = new_data.rename(columns=lambda x: x.replace('vol', 'volume'))
        self._market_data_day = self.merge_data(self._market_data_day, bar)

    def callback_week(self, a, b, c, body):
        new_data = pd.read_msgpack(body)
        bar = new_data.rename(columns=lambda x: x.replace('vol', 'volume'))
        self._market_data_week = self.merge_data(self._market_data_week, bar)

    def callback_month(self, a, b, c, body):
        new_data = pd.read_msgpack(body)
        bar = new_data.rename(columns=lambda x: x.replace('vol', 'volume'))
        self._market_data_month = self.merge_data(self._market_data_month, bar)

    def callback_year(self, a, b, c, body):
        new_data = pd.read_msgpack(body)
        bar = new_data.rename(columns=lambda x: x.replace('vol', 'volume'))
        self._market_data_year = self.merge_data(self._market_data_year, bar)

    def _debug_sim(self):
        self.running_mode = 'sim'

        for freq in self.frequences:
            if freq == '1min':
                self._market_data_1min = self.get_old_data(N=5, freq=freq)
            elif freq == '5min':
                self._market_data_5min = self.get_old_data(N=10, freq=freq)
            elif freq == '15min':
                self._market_data_15min = self.get_old_data(N=20, freq=freq)
            elif freq == '30min':
                self._market_data_30min = self.get_old_data(N=50, freq=freq)
            elif freq == '60min':
                self._market_data_60min = self.get_old_data(N=76, freq=freq)
            elif freq == 'day':
                self._market_data_day = self.get_old_data(N=300, freq=freq)
            elif freq == 'week':
                self._market_data_week = self.get_old_data(N=300, freq=freq)
            elif freq == 'month':
                self._market_data_month = self.get_old_data(N=500, freq=freq)
            elif freq == 'year':
                self._market_data_year = self.get_old_data(N=1000, freq=freq)

        self.database = pymongo.MongoClient(mongo_ip).QAREALTIME
        self.client = self.database.account
        self.subscriber_client = self.database.subscribe

        self.acc = QIFI_StockSIM_Account(username=self.username, password=self.password, init_cash=self.init_cash)
        self.acc.initial()

        self.pub = publisher_routing(exchange='QAORDER_ROUTER',
                                     host=self.trade_host, port=self.trade_port,
                                     user=self.trade_user, password=self.trade_password,
                                     routing_key=self.strategy_id)
        self.subscribe_datas(self.code, self.frequences, self.data_host,
                             self.data_port, self.data_user, self.data_password)

        self.database.strategy_schedule.job_control.update(
            {'strategy_id': self.strategy_id},
            {'strategy_id': self.strategy_id, 'taskid': self.taskid,
             'filepath': os.path.abspath(__file__), 'status': 200}, upsert=True)

    def get_old_data(self, N, freq):
        start = QA.QA_util_get_pre_trade_date(QA.QA_util_get_real_date(str(datetime.date.today())), n=N)
        _data = QA.QA_quotation(self.code, start=start, end=str(datetime.datetime.now()),
                                frequence=freq, market=self.market_type,
                                source=QA.DATASOURCE.MONGO, output=QA.OUTPUT_FORMAT.DATAFRAME)
        return _data.loc[:, ['open', 'high', 'low', 'close', 'volume', 'amount']]

    def get_freq_type(self, freq):
        if freq.endswith('min'):
            return QA.FREQUENCE.REALTIME_MIN
        elif freq.endswith('day'):
            return QA.FREQUENCE.DAY
        elif freq.endswith('week'):
            return QA.FREQUENCE.WEEK
        elif freq.endswith('month'):
            return QA.FREQUENCE.MONTH
        elif freq.endswith('year'):
            return QA.FREQUENCE.YEAR

    def run(self):
        while True:
            pass

    def run_sim(self):
        self._debug_sim()
        # self.sub_1min.start()
        for freq in self.frequences:
            if freq == '1min':
                # self.sub_1min.start()
                threading.Thread(target=self.sub_1min.start, daemon=True).start()
            elif freq == '5min':
                self.sub_5min.start()
            elif freq == '15min':
                self.sub_15min.start()
            elif freq == '30min':
                self.sub_30min.start()
            elif freq == '60min':
                self.sub_60min.start()
            elif freq == 'day':
                # self.sub_day.start()
                threading.Thread(target=self.sub_day.start, daemon=True).start()
            elif freq == 'week':
                self.sub_week.start()
            elif freq == 'month':
                self.sub_month.start()
            elif freq == 'year':
                self.sub_year.start()

    def debug(self):
        self.running_mode = 'backtest'
        self.database = pymongo.MongoClient(mongo_ip).QUANTAXIS
        user = QA_User(username='admin', password='admin')
        port = user.new_portfolio(self.portfolio)
        self.acc = port.new_accountpro(account_cookie=self.username,
                                       init_cash=self.init_cash,
                                       market_type=self.market_type,
                                       frequence=self.frequences)
        self.broker = QA.QA_BacktestBroker()

        data = QA.QA_quotation(self.code, self.start.strftime('%Y-%m-%d'), self.end.strftime('%Y-%m-%d'),
                               source=QA.DATASOURCE.MONGO, frequence='1min', market=self.market_type,
                               output=QA.OUTPUT_FORMAT.DATAFRAME)
        data.groupby(level=0, sort=False).apply(self.x1_1min)
        # for date in data.index.levels[0]:
        #     self.x1_1min(data.loc[(date, slice(None)), :])

    def x1(self, item):
        self.latest_price[item.name[1]] = item['close']
        if str(item.name[0])[0:10] != str(self.running_time)[0:10]:
            self.on_dailyclose()
            self.on_dailyopen()
            if self.market_type == QA.MARKET_TYPE.STOCK_CN or self.market_type == QA.MARKET_TYPE.INDEX_CN:
                self.acc.settle()
        # self._on_1min_bar()
        self._market_data.append(copy.deepcopy(item))
        self.running_time = str(item.name[0])
        self.on_bar_backtest(item)

    def x1_1min(self, item):
        # self.latest_price[item.name[1]] = item['close']
        if str(item.index.values[0][0])[0:10] != str(self.running_time)[0:10]:
            self.acc.settle()
            print('settle at %s' % str(item.index.values[0][0]))
        self._market_data_1min = self._market_data_1min.append(item)
        self.resample_market_data()
        self.running_time = str(item.index.values[0][0])
        self.on_bar_backtest(item)

    def resample_market_data(self):
        for freq in self.frequences:
            if freq != QA.FREQUENCE.ONE_MIN:
                if freq == '5min':
                    self._market_data_5min = self.market_data_1min \
                        .groupby(level=1, sort=False) \
                        .apply(lambda x: QA_data_min_resample(x.reset_index(1), type_='5min'))
                    self._market_data_5min.index = self._market_data_5min.index.droplevel(level=0)
                elif freq == '15min':
                    self._market_data_15min = self.market_data_1min \
                        .groupby(level=1, sort=False) \
                        .apply(lambda x: QA_data_min_resample(x.reset_index(1), type_='15min'))
                    self._market_data_15min.index = self._market_data_15min.index.droplevel(level=0)
                elif freq == '30min':
                    self._market_data_30min = self.market_data_1min \
                        .groupby(level=1, sort=False) \
                        .apply(lambda x: QA_data_min_resample(x.reset_index(1), type_='30min'))
                    self._market_data_30min.index = self._market_data_30min.index.droplevel(level=0)
                elif freq == '60min':
                    self._market_data_60min = self.market_data_1min \
                        .groupby(level=1, sort=False) \
                        .apply(lambda x: QA_data_min_resample(x.reset_index(1), type_='60min'))
                    self._market_data_60min.index = self._market_data_60min.index.droplevel(level=0)
                elif freq == 'day':
                    self._market_data_day = self.market_data_1min \
                        .groupby(level=1, sort=False) \
                        .apply(lambda x: QA_data_min_to_day(x.reset_index(1), type_='D'))
                elif freq == 'week':
                    self._market_data_week = self.market_data_day \
                        .groupby(level=1, sort=False) \
                        .apply(lambda x: QA_data_min_to_day(x.reset_index(1), type_='1W'))
                elif freq == 'month':
                    self._market_data_month = self.market_data_day \
                        .groupby(level=1, sort=False) \
                        .apply(lambda x: QA_data_min_to_day(x.reset_index(1), type_='1M'))
                elif freq == 'year':
                    self._market_data_year = self.market_data_day \
                        .groupby(level=1, sort=False) \
                        .apply(lambda x: QA_data_min_to_day(x.reset_index(1), type_='1Y'))

    def update_account(self):
        if self.running_mode == 'sim':
            self.accounts = self.acc.account_msg
            self.orders = self.acc.orders
            self.positions = self.acc.positions
            self.trades = self.acc.trades
            self.updatetime = self.acc.dtstr
        elif self.running_mode == 'backtest':
            self.positions = self.acc.positions

    def send_order(self,  direction='BUY', offset='OPEN', code=None, price=3925, volume=10, order_id='',):
        # print('start send order')
        towards = eval('ORDER_DIRECTION.{}_{}'.format(direction, offset))
        order_id = str(uuid.uuid4()) if order_id == '' else order_id

        if self.market_type == QA.MARKET_TYPE.STOCK_CN or self.market_type == QA.MARKET_TYPE.INDEX_CN:
            """
            在此对于股票的部分做一些转换
            """
            if towards == ORDER_DIRECTION.SELL_CLOSE:
                towards = ORDER_DIRECTION.SELL
            elif towards == ORDER_DIRECTION.BUY_OPEN:
                towards = ORDER_DIRECTION.BUY

        if isinstance(price, float):
            pass
        elif isinstance(price, pd.Series):
            price = price.values[0]

        if self.running_mode == 'sim':

            QA.QA_util_log_info(
                '============ {} SEND ORDER =================='.format(order_id))
            QA.QA_util_log_info('direction:{} offset:{} price:{} volume:{}'.format(
                direction, offset, price, volume))

            if self.check_order(direction, offset):
                self.last_order_towards = {'BUY': '', 'SELL': ''}
                self.last_order_towards[direction] = offset
                # now = str(datetime.datetime.now())
                # print(now)

                # 不在策略中更新账户，通过pub通知账户更新order macong 2020.08.27
                # order = self.acc.send_order(code=code, towards=towards, price=price, amount=volume, order_id=order_id)

                # print('acc make deal')
                # self.acc.make_deal(order)
                # self.bar_order['{}_{}'.format(direction, offset)] = self.bar_id

                order = {}
                order['topic'] = 'send_order'
                order['instrument_id'] = code
                order['direction'] = direction
                order['volume'] = volume
                order['price'] = price
                self.pub.pub(json.dumps(order), routing_key=self.strategy_id)
                # if self.send_wx:
                #     for user in self.subscriber_list:
                #         QA.QA_util_log_info(self.subscriber_list)
                #         try:
                #             "oL-C4w2WlfyZ1vHSAHLXb2gvqiMI"
                #             """http://www.yutiansut.com/signal?user_id=oL-C4w1HjuPRqTIRcZUyYR0QcLzo&template=xiadan_report&\
                #                         strategy_id=test1&realaccount=133496&code=rb1910&order_direction=BUY&\
                #                         order_offset=OPEN&price=3600&volume=1&order_time=20190909
                #             """
                #
                #             requests.post('http://www.yutiansut.com/signal?user_id={}&template={}&strategy_id={}&realaccount={}&code={}&order_direction={}&order_offset={}&price={}&volume={}&order_time={}'.format(
                #                 user, "xiadan_report", self.strategy_id, self.acc.user_id, code, direction, offset, price, volume, now))
                #         except Exception as e:
                #             QA.QA_util_log_info(e)

            else:
                QA.QA_util_log_info('failed in ORDER_CHECK')

        elif self.running_mode == 'backtest':

            self.bar_order['{}_{}'.format(direction, offset)] = self.bar_id

            self.acc.receive_simpledeal(
                code=code, trade_time=self.running_time, trade_towards=towards, trade_amount=volume, trade_price=price, order_id=order_id)
            #self.positions = self.acc.get_position(self.code)


if __name__ == '__main__':
    # QAStrategyStockBase(code=['000001', '000002']).run_sim()
    QAStrategyStockBase(username='sim_test1', password='123456', code=['002664'], start='2020-01-01', end='2020-05-21').run_backtest()
